#include "Common7.h"



TraCommMsgCallback g_TraCommMsgCallback = NULL;

FileCallBack OldExpDisQueryAttributeInformation = NULL;
FileCallBack OldExpDisSetAttributeInformation = NULL;
PULONG64 gWin7ExpDisQueryAttributeInformation = NULL;



NTSTATUS MyCommonFuncSetFileForWin7(HANDLE FileHandle, PVOID info, PVOID param3, PVOID param4)
{

	//DbgPrintEx(77, 0, "into MyCommonFuncSetFileForWin7 func\r\n");
	if (MmIsAddressValid(info))
	{
		stCommonDate* pMsgInfo = (stCommonDate*)info;
		NTSTATUS status = g_TraCommMsgCallback(pMsgInfo);
		
		if (!NT_SUCCESS(status))
		{
			if (OldExpDisSetAttributeInformation)
			{
				return OldExpDisSetAttributeInformation(FileHandle, info, param3, param4);
			}
		}
		
		return STATUS_SUCCESS;
	}


}

NTSTATUS MyCommonFuncQueryFileForWin7(HANDLE FileHandle, PVOID info, PVOID param3, PVOID param4)
{

	//DbgPrintEx(77, 0, "into MyCommonFuncSetFileForWin7 func\r\n");
	if (MmIsAddressValid(info))
	{
		stCommonDate* pMsgInfo = (stCommonDate*)info;
		NTSTATUS status = g_TraCommMsgCallback(pMsgInfo);

		if (!NT_SUCCESS(status))
		{
			if (OldExpDisQueryAttributeInformation)
			{
				return OldExpDisQueryAttributeInformation(FileHandle, info, param3, param4);
			}
		}

		return STATUS_SUCCESS;
	}

}

NTSTATUS InstallCommonForWIn7(TraCommMsgCallback TraMsgCallBack)
{

	UNICODE_STRING strExQueryAttributeInformation = { 0 };
	RtlInitUnicodeString(&strExQueryAttributeInformation, L"ExRegisterAttributeInformationCallback");

	PVOID pFuncAddr = MmGetSystemRoutineAddress(&strExQueryAttributeInformation);

	if (pFuncAddr == NULL) return STATUS_FAIL_CHECK;

	ULONG64 offset = *(PLONG)((ULONG64)pFuncAddr + 0xd + 3);


	PULONG64 ExpDisQueryAttributeInformation_Addr = (PULONG64)((ULONG64)pFuncAddr + 0xd + 7 + offset);

	OldExpDisQueryAttributeInformation = (FileCallBack)(ExpDisQueryAttributeInformation_Addr[0]);

	OldExpDisSetAttributeInformation = (FileCallBack)(ExpDisQueryAttributeInformation_Addr[1]);

	RegisterFileCallBack RegCallbackFunc = (RegisterFileCallBack)pFuncAddr;

	ExpDisQueryAttributeInformation_Addr[0] = 0;
	ExpDisQueryAttributeInformation_Addr[1] = 0;

	stRegisterFileCallBack fileCallback = { 0 };
	fileCallback.QueryFileCallBack = MyCommonFuncQueryFileForWin7;
	fileCallback.SetFileCallBack = MyCommonFuncSetFileForWin7;

	NTSTATUS status = RegCallbackFunc(&fileCallback);

	if (!NT_SUCCESS(status))
	{
		DbgPrintEx(77, 0, "[wz]Install Common For Win7 Error :%d\r\n", status);
		return status;
	}

	g_TraCommMsgCallback = TraMsgCallBack;
	gWin7ExpDisQueryAttributeInformation = ExpDisQueryAttributeInformation_Addr;
	DbgPrintEx(77, 0, "[wz]Install Common For Win7 Sucess !\r\n");

	return status;
}

NTSTATUS UninstallCommonForWIn7()
{
	if (gWin7ExpDisQueryAttributeInformation)
	{
		gWin7ExpDisQueryAttributeInformation[0] = OldExpDisQueryAttributeInformation;
		gWin7ExpDisQueryAttributeInformation[1] = OldExpDisSetAttributeInformation;
	}
	DbgPrintEx(77, 0, "[wz]Unitall Common For Win7 Sucess !\r\n");
	return STATUS_SUCCESS;
}
